// vite-plugin-inject-context-name.js
import generate from '@babel/generator'
import * as parser from '@babel/parser'
import traverse from '@babel/traverse'
import * as t from '@babel/types'
import type { Plugin } from 'vite'

export default function InjectContextNamePlugin(): Plugin {
  return {
    name: 'vite-plugin-inject-context-name',
    transform(code, id) {
      if ((id.endsWith('.ts') || id.endsWith('.tsx'))) {
        const ast = parser.parse(code, {
          sourceType: 'module',
          plugins: [
            'typescript',
            'jsx',
          ],
        })

        let isVueContext = false // 假设首先不是Vue的 createContext

        traverse(ast, {
          // 查找 createContext 是否来自 React
          ImportDeclaration(path) {
            if (path.node.source.value === 'react') {
              const specifiers = path.node.specifiers
              specifiers.forEach((spec) => {
                if (t.isImportSpecifier(spec))
                  console.log(spec.imported)

                if (
                  t.isImportSpecifier(spec)
                  && spec.imported.name === 'createContext'
                ) {
                  isVueContext = false
                  path.stop() // 找到了就停止遍历
                }
              })
            }
          },
          // 如果 createContext 没有来自 React，设置标记为 Vue
          CallExpression(path) {
            if (
              t.isIdentifier(path.node.callee)
              && path.node.callee.name === 'createContext'
            )
              isVueContext = true
          },
          VariableDeclarator(path) {
            if (isVueContext) {
              const node = path.node
              if (
                t.isCallExpression(node.init)
                && t.isIdentifier(node.init.callee)
                && node.init.callee.name === 'createContext'
                && t.isIdentifier(node.id)
              ) {
                const contextName = node.id.name // The variable name
                if (node.init.arguments.length === 1) {
                  // Modify the AST to inject the context name
                  node.init.arguments.push(t.arrowFunctionExpression([], t.blockStatement([])))
                }

                // Modify the AST to inject the context name
                node.init.arguments.push(t.stringLiteral(`${contextName}.Provider`))
              }
            }
          },
        })

        if (isVueContext) {
          // 如果确定是Vue的 createContext，转换AST回代码
          const { code: transformedCode } = generate(ast)
          return {
            code: transformedCode,
            map: null, // 需要的话这里应该返回正确的sourcemap
          }
        }
      }
      // 如果不需要转换，则直接返回
      return {
        code,
        map: null,
      }
    },
  }
}
